#include "kernel_operator.h"
using namespace AscendC;

template<typename TYPE_X,typename TYPE_INDICES,typename TYPE_AXIS>
class KernelGatherV3
{
public:
    __aicore__ inline KernelGatherV3(){}//构造函数
    __aicore__ inline uint64_t compute_dims_product(uint64_t *shape, uint64_t start, uint64_t end ){
            if(start>=end) return 1;
            int product = 1;
            for(int i = start ; i < end ; i++ ){
                product*=shape[i];
            }
            return product;
        }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indices,GM_ADDR axisa,GM_ADDR out,//输入
                                uint64_t batch_dims,uint64_t x_dim,uint64_t indices_dim,
                                uint64_t * x_ndarray,uint64_t * indices_ndarray,uint64_t * out_ndarray)
    {   
            
        axisGm.SetGlobalBuffer((__gm__ TYPE_AXIS *)axisa);
        this -> axis = axisGm.GetValue(0);

        this -> x_ndarray = x_ndarray;
        this -> indices_ndarray = indices_ndarray;

        this -> x_dim = x_dim;
        this ->indices_dim =indices_dim;
        this->batch_dims=batch_dims;
        this ->axis_dim = x_ndarray[axis];



        this -> batch_size =  compute_dims_product(x_ndarray,0,batch_dims);//批处理大小
        this -> outer_dims =  compute_dims_product(x_ndarray,batch_dims,axis);//轴前维度
        this -> indices_inner_count = compute_dims_product(indices_ndarray,batch_dims,indices_dim);//计算索引维度
        this -> inner_dims = compute_dims_product(x_ndarray,axis+1,x_dim);//轴后维度
        this -> out_ndarray = out_ndarray;
        xGm.SetGlobalBuffer((__gm__ TYPE_X *)x);
        indicesGm.SetGlobalBuffer((__gm__ TYPE_INDICES *)indices);
        outGm.SetGlobalBuffer((__gm__ TYPE_X*)out);
    
    }
    __aicore__ inline void Process()
    {   
        for(uint64_t i = 0 ;i<batch_size ;i++)
            {
                compute(i);
            }
    }
private:
    __aicore__ inline void compute(uint64_t progress)
    {
        uint64_t params_size_per_batch = outer_dims * axis_dim * inner_dims;
        uint64_t out_size_per_batch = outer_dims * indices_inner_count * inner_dims;
        uint64_t indices_size_per_batch = indices_inner_count;
        //不同batch偏移
        uint64_t params_batch_offset = progress * params_size_per_batch;
        uint64_t indices_batch_offset = progress * indices_size_per_batch;
        uint64_t output_batch_offset = progress * out_size_per_batch;

        //在当前batch中进行gather
        for(int outer_idx =0 ;outer_idx < outer_dims;outer_idx ++){
            for(int idx = 0 ;idx < indices_inner_count ;idx++){
                int slice_idx = indicesGm.GetValue(indices_batch_offset +idx);
                //内部循环
                for(int inner_idx =0 ; inner_idx < inner_dims;inner_idx++){
                    int input_pos = params_batch_offset + 
                                  (outer_idx * axis_dim + slice_idx) * inner_dims + 
                                  inner_idx;
                    
                    // 计算输出位置
                    int output_pos = output_batch_offset + 
                                   (outer_idx * indices_inner_count + idx) * inner_dims + 
                                   inner_idx;
                    
                    // 复制数据
                    outGm.SetValue(output_pos,xGm.GetValue(input_pos));
                }
            }
        }
    }
private:
    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_INDICES> indicesGm;
    GlobalTensor<TYPE_X> outGm;
    GlobalTensor<TYPE_AXIS> axisGm;
    uint64_t batch_size;
    uint64_t batch_dims;
    uint64_t outer_dims;
    uint64_t axis;
    uint64_t indices_dim;
    uint64_t axis_dim;
    uint64_t x_dim;
    uint64_t indices_inner_count;
    uint64_t inner_dims;
    uint64_t* x_ndarray;
    uint64_t* indices_ndarray;
    uint64_t* out_ndarray;
};

extern "C" __global__ __aicore__ void gather_v3(GM_ADDR x, GM_ADDR indices, GM_ADDR axis, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelGatherV3<DTYPE_X,DTYPE_INDICES,DTYPE_AXIS> op;
    op.Init(x,indices,axis,out,tiling_data.batch_dims,tiling_data.x_dim,tiling_data.indices_dim,
                               tiling_data.x_ndarray,tiling_data.indices_ndarray,tiling_data.out_ndarray);
    op.Process();
}